#!/usr/bin/env python2
# -*- coding:utf-8 -*-

import codecs, hashlib, base64
import os, sys
import re
from datetime import datetime

UnKnownUser = "UnKnownUser"

class AuthorizedKey:
  def __init__(self, key_string):
    self.load(key_string)
  def load(self, key_string):
    arr = key_string.split(' ') + [UnKnownUser]
    key = arr[1].strip()
    self.user = arr[2].strip()
    self.md5 = self.md5(key)
    self.sha256 = self.sha256(key)

  def md5(self, key):
    fp = hashlib.md5(codecs.decode(bytes(key), 'base64')).hexdigest()
    return ":".join([fp[pair:pair+2] for pair in range(0, len(fp)-1, 2) ])
  
  def sha256(self, key):
    sha256 = hashlib.sha256()
    sha256.update(base64.b64decode(key))
    b64fingerprint = base64.b64encode(sha256.digest())
    # ssh-keygen会去掉 '='
    return b64fingerprint.decode("utf-8").replace("=", '')

  def puts(self):
    print("{0} {1} {2}".format(self.user, self.md5, self.sha256))

class KeyManager:
  def __init__(self, *paths):
    self.md5_map = {}
    self.sha256_map = {}
    if len(paths) == 0:
      paths = ["/root", "/home/*"]
    self.load_paths(paths)

  def load_paths(self, paths):
    self.paths = []
    for path in paths:
      if path.endswith("*"):
        new_path = path.replace("*", "")
        for dir in os.listdir(new_path):
          ak_path = os.path.join(new_path, dir, ".ssh/authorized_keys")
          self.paths.append(ak_path)
      elif os.path.isdir(path):
        ak_path = os.path.join(path, ".ssh/authorized_keys")
        self.paths.append(ak_path)
      else:
        self.paths.append(path)
  def load_authorized_keys(self):
    for path in self.paths:
      self.load_path(path)
    
  def load_path(self, path):
    try:
      with open(path, 'r') as f:
        for line in f:
          if self.not_valid(line):
            continue
          ak = AuthorizedKey(line)
          if self.md5_map.get(ak.md5) == None or self.md5_map[ak.md5].user == UnKnownUser:
            self.md5_map[ak.md5] = ak
          if self.sha256_map.get(ak.sha256) == None or self.sha256_map[ak.sha256].user == UnKnownUser:
            self.sha256_map[ak.sha256] = ak
    except IOError as err:
      print("open {0} {1}".format(path, err))
  
  def not_valid(self, line):
    if "ssh-rsa" in line:
      return False
    return True

  def find(self, hash_str):
    if ":" in hash_str:
      return self.md5_map.get(hash_str)
    else:
      return self.sha256_map.get(hash_str)

  def puts(self):
    for k,v in self.md5_map.items():
      print(k, v.user)

class LoginParser:
  def __init__(self, key_manager, path):
    self.key_manager = key_manager
    self.path = path
    if not os.path.exists(path):
      print("LoginParser ERROR not exists {0}".format(path))
      return None
    print("LoginParser INFO from {0}".format(path))
    if "auth.log" in path:
      self.auth_log()
    elif "audit.log" in path:
      self.audit_log()
    elif "secure" in path:
      self.secure_log()

  # 分析ubuntu的 auth.log reg可以在 https://pythex.org/ 测试使用
  AUTH_LOG_ACCEPTED_REG = "^(.+?\d{1,2}\s\d{1,2}:\d{1,2}:\d{1,2})\s.*\sAccepted publickey for (.+) from (.*) port .* SHA256:(.*)$"
  def auth_log(self):
    with open(self.path, 'rb') as log:
      for line in log:
        if "Accepted publickey" in line:
          m = re.search(self.AUTH_LOG_ACCEPTED_REG, line)
          if m != None:
            ak = self.key_manager.find(m.group(4))
            le = LoginEvent(time = m.group(1), ruser = m.group(2), user=UnKnownUser, fip=m.group(3), hash_str=m.group(4), ak=ak)
            le.puts()
  # 分析 audit.log reg可以在 https://pythex.org/ 测试使用
  AUDIT_LOG_USER_AUTH_REG = 'USER_AUTH.*msg=audit\((\d+)\.\d+\:\d+\):.*fp=(.*) rport.*acct="(\w*)".*addr=(.*) terminal'
  def audit_log(self):
    with open(self.path, 'rb') as log:
      for line in log:
        m = re.search(self.AUDIT_LOG_USER_AUTH_REG, line)
        if m != None:
          ak = self.key_manager.find(m.group(2))
          le = LoginEvent(time = m.group(1), ruser = m.group(3), user=UnKnownUser, fip=m.group(4), hash_str=m.group(2), ak=ak)
          le.puts()
  # 分析 secure reg可以在 https://pythex.org/ 测试使用
  SECURE_FOUND_MATCHING_REG = '(.+?\d{1,2}\s\d{1,2}:\d{1,2}:\d{1,2})\s.*Found matching RSA key:\s(.+?)\s'
  SECURE_ACCEPTED_PUBLICKEY_REG = '(.+?\d{1,2}\s\d{1,2}:\d{1,2}:\d{1,2})\s.*Accepted publickey for\s(.+?)\sfrom\s(.+?)\s'
  def secure_log(self):
    ak = None
    with open(self.path, 'rb') as log:
      for line in log:
        m = re.search(self.SECURE_FOUND_MATCHING_REG, line)
        if m != None:
          ak = self.key_manager.find(m.group(2))
          le = LoginEvent(time = m.group(1), ruser = "", user=UnKnownUser, fip="", hash_str=m.group(2), ak=ak)
          le.puts()
          continue
        m = re.search(self.SECURE_ACCEPTED_PUBLICKEY_REG, line)
        if m != None:
          le = LoginEvent(time = m.group(1), ruser = m.group(2), user=UnKnownUser, fip=m.group(3), hash_str="", ak=ak)
          le.puts()
          ak = None

class LoginEvent:
  def __init__(self, time, ruser, user, hash_str, fip, ak):
    self.time = self.timeformat(time)
    self.ruser = ruser
    self.user = user
    self.hash_str = hash_str
    self.fip = fip
    if ak != None:
      self.user = ak.user
      self.hash_str = ak.md5
  def timeformat(self, time):
    try:
      if re.match("\d+?$", time):
        return "{0}".format(datetime.fromtimestamp(int(time)))
      elif re.match("(.+?\s\d{1,2}\s\d{1,2}:\d{1,2}:\d{1,2})$", time):
        # https://www.delftstack.com/zh/howto/python/how-to-convert-string-to-datetime/
        dt = datetime.strptime("{0} {1}".format(datetime.today().year, time), "%Y %b %d %H:%M:%S")
        return "{0}".format(dt)
      else:
        return time
    except IOError as err:
      print("format time {0} ERROR {1}".format(time, err))
      return time

  def puts(self):
    print(",".join([self.time, self.ruser, self.user, self.hash_str, self.fip]))

if __name__ == '__main__':
  print("Run with python", "".join(sys.version.split("\n")))
  km = KeyManager()
  km.load_authorized_keys()
  km.puts()
  LoginParser(km, "/var/log/auth.log")
  LoginParser(km, "/var/log/audit/audit.log")
  LoginParser(km, "/var/log/secure")
